(* Compile: Convert predefined functions into bytecode series. *)
open Ast
open Bytecode

module StringMap = Map.Make(String)

(* Symbol table: Information about all the names in scope *)
type env = {
    function_index : int StringMap.t; (* Index for each function *)
    global_index   : int StringMap.t; (* "Address" for global variables *)
    local_index    : int StringMap.t; (* FP offset for args, locals *)
  }

(* val enum : int -> 'a list -> (int * 'a) list *)
let rec enum stride n = function (*Change here: this func maps variable name to offsets. *)
    [] -> []
  | hd::tl ->  (* hd is the vari_decl  *)
	if stride > 0 then
		match hd.vtype with
			"int" -> (* Allocate global storage space for an int *)
				(n + 1, hd.name) :: enum stride (n+stride * 2) tl
 		|	"string" -> (* Here's the question: how many slots need to be allocated to string? *)
						(* To make it simpler, allocate 30 slots for it *)
				(n + 29, hd.name) :: enum stride (n+stride * 30) tl 
 		|	"line" -> 
				(n + 4, hd.name) :: enum stride (n+stride * 5) tl 
 		|	"rect" ->
				(n + 4, hd.name) :: enum stride (n+stride * 5) tl 
 		|	"ellipse" -> 
				(n + 4, hd.name) :: enum stride (n+stride * 5) tl 
 		|	"shape" -> 
				(n + 4, hd.name) :: enum stride (n+stride * 5) tl 
        |   "arrayi" ->
				(n + 2*hd.vsize-1, hd.name) :: enum stride (n+stride * 2 * hd.vsize) tl
        |   "arrayr" ->
				(n + 5*hd.vsize-1, hd.name) :: enum stride (n+stride * 5 * hd.vsize) tl
        |   "arraye" ->
				(n + 5*hd.vsize-1, hd.name) :: enum stride (n+stride * 5 * hd.vsize) tl
        |   "arrayl" ->
				(n + 5*hd.vsize-1, hd.name) :: enum stride (n+stride * 5 * hd.vsize) tl
        |   "arrays" ->
				(n + 5*hd.vsize-1, hd.name) :: enum stride (n+stride * 5 * hd.vsize) tl
 		|	"bind" -> 
						(* so the problem is with bind: allocate 100 slots. *)
				(n + 99, hd.name) :: enum stride (n+stride * 100) tl 
		| 	_ -> raise(Failure ("Undefined type with variable" ^ hd.name))
	else
		match hd.vtype with
			"int" -> (* Allocate global storage space for an int *)
				(n, hd.name) :: enum stride (n+stride * 2) tl
 		|	"string" -> (* Here's the question: how many slots need to be allocated to string? *)
						(* To make it simpler, allocate 30 slots for it *)
				(n, hd.name) :: enum stride (n+stride * 30) tl 
 		|	"line" -> 
				(n, hd.name) :: enum stride (n+stride * 5) tl 
 		|	"rect" ->
				(n, hd.name) :: enum stride (n+stride * 5) tl 
 		|	"ellipse" -> 
				(n, hd.name) :: enum stride (n+stride * 5) tl 
 		|	"shape" -> 
				(n, hd.name) :: enum stride (n+stride * 5) tl 
 		|	"arrayi" -> 
				(n, hd.name) :: enum stride (n+stride * 2 * hd.vsize) tl 
 		|	"arrayr" -> 
				(n, hd.name) :: enum stride (n+stride * 5 * hd.vsize) tl 
 		|	"arraye" -> 
				(n, hd.name) :: enum stride (n+stride * 5 * hd.vsize) tl 
 		|	"arrayl" -> 
				(n, hd.name) :: enum stride (n+stride * 5 * hd.vsize) tl 
 		|	"arrays" -> 
				(n, hd.name) :: enum stride (n+stride * 5 * hd.vsize) tl 
		| 	_ -> raise(Failure ("Undefined type with variable " ^ hd.name))

(* val enum : int -> 'a list -> (int * 'a) list *)
let rec enum_func stride n = function
    [] -> []
  | hd::tl -> (n, hd) :: enum_func stride (n+stride) tl

let get_vari_size a vlist = 
	 List.fold_left (fun a b -> a + (match b.vtype with
										"int" -> 2	
									|   "string" -> 30
									|   "rect" -> 5
									|   "ellipse" -> 5
									|   "line" -> 5
									|   "shape" -> 5
									|   "bind"  -> 100 
									| 	"arrayi" -> b.vsize*2
									| 	"arrayr" -> b.vsize*5
									| 	"arraye" -> b.vsize*5
									| 	"arrayl" -> b.vsize*5
									| 	"arrays" -> b.vsize*5
									|   _ -> raise(Failure("Error in get_vari_size !!"))
									)) 0 vlist

(* val string_map_pairs StringMap 'a -> (int * 'a) list -> StringMap 'a *)
let string_map_pairs map pairs =
  List.fold_left (fun m (i, n) -> StringMap.add n i m) map pairs (* Add (var_name, offset) pairs *)

(** Translate a program in AST form into a bytecode program.  Throw an
    exception if something is wrong, e.g., a reference to an unknown
    variable or function *)
(* The input of translate are two lists: one for global, one for function. Compiler 
   takes the two lists, which is generated by parser, to be bytecode stuff. *)
let translate (globals, functions) = (* globals has the form of vari_decl *)

  (* Allocate "addresses" and storing place for each global variable *)
  let global_indexes = string_map_pairs StringMap.empty (enum 1 0 globals) in

  (* Assign indexes to function names; built-ins are special *)
  let built_in_functions = StringMap.add "print" (-1) StringMap.empty in
  let built_in_functions = StringMap.add "halt" (-2) built_in_functions in
	(* -3 maps to MakeBind which is deleted *)
  let built_in_functions = StringMap.add "printarray" (-3) built_in_functions in
  let built_in_functions = StringMap.add "print_int" (-4) built_in_functions in
  let built_in_functions = StringMap.add "dumpstack" (-5) built_in_functions in
  let function_indexes = string_map_pairs built_in_functions
      (enum_func 1 1 (List.map (fun f -> f.fname) functions)) in

  (* Translate a function in AST form into a list of bytecode statements *)
  let translate env fdecl =
    (* Bookkeeping: FP offsets for locals and arguments *)
    let size_formals = get_vari_size 0 fdecl.formals 
    and size_locals = get_vari_size 0 fdecl.locals 
    and local_offsets = enum 1 1 fdecl.locals (* This is list of pair (#, local_name) *)
    and formal_offsets = enum (-1) (-2) fdecl.formals in (* The 1 and -2 in this and last line means the first variable's address, if any. *)
    let env = { env with local_index = string_map_pairs (* Not quite understand this line. *)
		  StringMap.empty (local_offsets @ formal_offsets) } in

    let rec expr = function (* Try to implement built-in functions here. *)
	Literal i -> [Litin i] (* Give input as an int. *) 
      | String s -> (let s2 = String.sub s 3 (String.length s - 4)
		in [Litst s2])
      | Rectangle r -> let r2 = String.sub r 3 (String.length r - 4) (* No need of type check. Done in parser. *)
		in [Litsh ("3 " ^ r2)]
      | Ellipse e -> let e2 = String.sub e 3 (String.length e - 4)
		in [Litsh ("4 " ^ e2)]
      | Line l -> let l2 = String.sub l 3 (String.length l - 4)
		in [Litsh ("5 " ^ l2)]
      | Shape s -> let s2 = String.sub s 3 (String.length s - 4)
		in [Litsh ("6 " ^ s2)]
      | Id s -> (* fetch variables according to id. put this on top of stack. *) 
	  (try [Lfp (StringMap.find s env.local_index)]
          with Not_found -> try [Lod (StringMap.find s env.global_index)]
          with Not_found -> raise (Failure ("undeclared variable " ^ s)))
      | Binop (e1, op, e2) -> expr e1 @ expr e2 @ [Bin op] (* Take care + might be overloaded to binds. *)
      | Assign (s, e) -> expr e @
	  (try [Sfp (StringMap.find s env.local_index)]
  	  with Not_found -> try [Str (StringMap.find s env.global_index)]
	  with Not_found -> raise (Failure ("undeclared variable " ^ s)))
      | Call (fname, actuals) -> (try (* How about the special functions about opengraph, close graph, halt,? *)
	  (List.concat (List.map expr (List.rev actuals))) @
	  [ Jsr (StringMap.find fname env.function_index) ]   
        with Not_found -> raise (Failure ("undefined function " ^ fname)))
      | Noexpr -> []
	  | Moveby(id, e1, e2) -> (* Move only handles with numbers*)
		(try [Lfp (StringMap.find id env.local_index)] 
			with Not_found -> try[Lod (StringMap.find id env.global_index)]
			with Not_found -> raise (Failure ("undeclared variable " ^ id)))
		@ expr e2 @ expr e1 @ [Movby] @ [Drp] @ [Drp] @
		(try [Sfp (StringMap.find id env.local_index)] 
			with Not_found -> try[Str (StringMap.find id env.global_index)]
			with Not_found -> raise (Failure ("undeclared variable " ^ id)))

	  | Moveto(id, e1, e2) -> (* Move only handles with numbers*)
		(try [Lfp (StringMap.find id env.local_index)] 
			with Not_found -> try[Lod (StringMap.find id env.global_index)]
			with Not_found -> raise (Failure ("undeclared variable " ^ id)))
		@ expr e2 @ expr e1 @ [Movto] @ [Drp] @ [Drp] @
		(try [Sfp (StringMap.find id env.local_index)] 
			with Not_found -> try[Str (StringMap.find id env.global_index)]
			with Not_found -> raise (Failure ("undeclared variable " ^ id)))

	  | GetV1(id) -> (* Get first value of object *)
		(try [Lfp (StringMap.find id env.local_index)]
			with Not_found -> try[Lod (StringMap.find id env.global_index)]
			with Not_found -> raise (Failure ("undeclared variable " ^ id)))
		@ [Litin 1] @ [GetC]

	  | GetV2(id) -> (* Get second value of object *)
		(try [Lfp (StringMap.find id env.local_index)]
			with Not_found -> try[Lod (StringMap.find id env.global_index)]
			with Not_found -> raise (Failure ("undeclared variable " ^ id)))
		@ [Litin 2] @ [GetC]
	
	  | GetV3(id) -> (* Get third value of object *)
		(try [Lfp (StringMap.find id env.local_index)]
			with Not_found -> try[Lod (StringMap.find id env.global_index)]
			with Not_found -> raise (Failure ("undeclared variable " ^ id)))
		@ [Litin 3] @ [GetC]

	  | GetV4(id) -> (* Get fourth value of object *)
		(try [Lfp (StringMap.find id env.local_index)]
			with Not_found -> try[Lod (StringMap.find id env.global_index)]
			with Not_found -> raise (Failure ("undeclared variable " ^ id)))
		@ [Litin 4] @ [GetC]
	  
	  | Aid(id, e) -> (* Retrieve Array value on top of stack *)
		expr e @ (try [Lfpa(StringMap.find id env.local_index)]
			with Not_found -> try[Loda (StringMap.find id env.global_index)]
			with Not_found -> raise (Failure ("undeclared variable" ^ id)))
		
	  | AAssign(id, e1, e2) -> (* Assign value to an array element. *)
		expr e2 @ expr e1 @ (try [Sfpa(StringMap.find id env.local_index)] 
			with Not_found -> try[Stra(StringMap.find id env.global_index)]
			with Not_found -> raise (Failure ("undelcared variable" ^ id)))

	  | REvaluation(v1, v2, v3, v4) -> (* Evaluate every variable in a rectangle*)
		expr v4 @ expr v3 @ expr v2 @ expr v1 @ [Litin 3] @ [MakeS]
		
	  | EEvaluation(v1, v2, v3, v4) -> (* Evaluate every variable in a rectangle*)
		expr v4 @ expr v3 @ expr v2 @ expr v1 @ [Litin 4] @ [MakeS]

	  | LEvaluation(v1, v2, v3, v4) -> (* Evaluate every variable in a rectangle*)
		expr v4 @ expr v3 @ expr v2 @ expr v1 @ [Litin 5] @ [MakeS]

    in let rec stmt = function
	Block sl     ->  List.concat (List.map stmt sl)
      | Expr e       -> expr e @ [Drp]
      | Return e     -> expr e @ [Rts size_formals]
      | If (p, t, f) -> let t' = stmt t and f' = stmt f in
	expr p @ [Beq(2 + List.length t')] @
	t' @ [Bra(1 + List.length f')] @ f'
      | For (e1, e2, e3, b) ->
	  stmt (Block([Expr(e1); While(e2, Block([b; Expr(e3)]))]))
      | While (e, b) ->
	  let b' = stmt b and e' = expr e in
	  [Bra (1+ List.length b')] @ b' @ e' @
	  [Bne (-(List.length b' + List.length e'))]

    in [Ent size_locals] @      (* Entry: allocate space for locals. Need to change as each variable has different size. *)
    stmt (Block fdecl.body) @  (* Body *)
    [Litin 0; Rts size_formals]   (* Default = return 0 *)

  in let env = { function_index = function_indexes;
		 global_index = global_indexes;
		 local_index = StringMap.empty } in

  (* Code executed to start the program: Jsr run; halt *)
  let entry_function = try
    [Sgraph; Jsr (StringMap.find "run" function_indexes); Hlt]
  with Not_found -> raise (Failure ("no \"run\" function"))
  in
    
  (* Compile the functions *)
  let func_bodies = entry_function :: List.map (translate env) functions in

  (* Calculate function entry points by adding their lengths *)
  let (fun_offset_list, _) = List.fold_left
      (fun (l,i) f -> (i :: l, (i + List.length f))) ([],0) func_bodies in
  let func_offset = Array.of_list (List.rev fun_offset_list) in

(* This is what compiler generates: *)
  { size_globals = get_vari_size 0 globals; (* This globals must be a complete layout of global var. *)
    (* Concatenate the compiled functions and replace the function
       indexes in Jsr statements with PC values *)
    text = Array.of_list (List.map (function
	Jsr i when i > 0 -> Jsr func_offset.(i)
      | _ as s -> s) (List.concat func_bodies))
  }
